package RSA;

import java.math.BigInteger;
import java.nio.charset.StandardCharsets;
import java.security.MessageDigest;
import java.security.NoSuchAlgorithmException;
import java.util.Random;

// 可以看看博客 https://blog.csdn.net/RedhatJongkhurun/article/details/124539180
public class OAEP_RSA {
    private final RSA RSA_CLASSIC = new RSA();

    private static final MessageDigest SHA1;

    private byte[] seed = null;

    private static final Random RANDOM = new Random();

    static {
        try {
            SHA1 = MessageDigest.getInstance("SHA-1");
        } catch (NoSuchAlgorithmException e) {
            throw new RuntimeException(e);
        }
    }

    public OAEP_RSA(byte[] seed) {
        if (seed.length != 20) {
            throw new RuntimeException("seed的长度应为160比特！");
        }
        this.seed = seed;
    }

    /**
     * OAEP的加密算法
     *
     * @param plaintext 待加密的消息
     * @param L         编码所需的标签，用以哈希算法的输入
     * @param k         编码的长度
     * @param e         公钥 (e,n) 中的 e
     * @param n         公钥 (e,n) 中的 n
     * @return 返回加密之后的一个大整数
     */
    public BigInteger encrypt(byte[] plaintext, String L, int k, BigInteger e, BigInteger n) {
        byte[] EM = encode(plaintext, L, k);
        BigInteger EMBigInt = new BigInteger(EM);
        return RSA_CLASSIC.encrypt(EMBigInt, e, n);
    }

    public byte[] decrypt(BigInteger ciphertext, String L, int k, BigInteger d, BigInteger n) {
        BigInteger plaintext = RSA_CLASSIC.decrypt(ciphertext, d, n);
        byte[] EM = plaintext.toByteArray();
        return decode(EM, L, k);
    }


    /**
     * 填充算法
     *
     * @param plainOAEP 输入的明文消息 Message
     * @param L         参数标签L，用以作为hash函数的输入，参数标签L应当来源于消息 Message
     * @param k         编码的长度
     * @return 填充编码后的消息EM，EM：Encoded Message
     */
    private byte[] encode(byte[] plainOAEP, String L, int k) {
        int mLen = plainOAEP.length;

        byte[] hash1 = SHA1.digest(L.getBytes(StandardCharsets.UTF_8));
        int hLen = hash1.length;

        int dbLen = k - 1 - hLen; // hLen其实是32
        // 检查长度，判断是否能对消息进行填充，DB的长度 ≥ 消息长度 + 哈希值长度(32字节) + 1字节(0x01)
        if (dbLen < mLen + hLen + 1) {
            throw new RuntimeException("编码长度太小！");
        }
        // 计算 PS(PaddingString) 的长度
        int psLen = dbLen - (hLen + 1 + mLen);

        // 拼接 DB = hash1 || PS || 0x01 || plainOAEP
        byte[] DB = new byte[dbLen];
        System.arraycopy(hash1, 0, DB, 0, hLen);
        for (int i = 0; i < psLen; i++) {
            DB[i + hLen] = 0x00;
        }
        DB[hLen + psLen] = 0x01;
        System.arraycopy(plainOAEP, 0, DB, 1 + hLen + psLen, mLen);

        // seed的长度固定为 hLen，

        // 计算 maskedDB
        byte[] DBmask = MGF(seed, dbLen);
        byte[] maskedDB = new byte[dbLen];
        for (int i = 0; i < dbLen; i++) {
            assert DBmask != null;
            maskedDB[i] = (byte) (DB[i] ^ DBmask[i]);
        }

        // 计算 seedmask
        byte[] seedmask = MGF(maskedDB, hLen);
        // 计算 maskedSeed
        byte[] maskedSeed = new byte[hLen];
        for (int i = 0; i < hLen; i++) {
            assert seedmask != null;
            maskedSeed[i] = (byte) (seed[i] ^ seedmask[i]);
        }

        byte[] res = new byte[k];
        res[0] = 0x01;
        System.arraycopy(maskedSeed, 0, res, 1, hLen);
        System.arraycopy(maskedDB, 0, res, 1 + hLen, dbLen);

        return res;
    }

    private byte[] decode(byte[] EM, String L, int k) {
        // 检查长度是否正确
        if (EM.length != k) {
            throw new RuntimeException("EM长度和编码长度k值不相等！");
        }

        byte[] hash1 = SHA1.digest(L.getBytes(StandardCharsets.UTF_8));
        int hLen = hash1.length;
        // 根据长度将EM分为 Y、maskedSeed、maskedDB
        // Y的长度为1，maskedSeed的长度为hLen，maskedDB长度为k-hLen-1，记为dbLen
        int dbLen = k - hLen - 1;

        byte Y = EM[0];
        if (Y != 0x01) {
            throw new RuntimeException("解码时第一个字节不是0x01！");
        }

        byte[] maskedSeed = new byte[hLen];
        System.arraycopy(EM, 1, maskedSeed, 0, hLen);

        byte[] maskedDB = new byte[dbLen];
        System.arraycopy(EM, 1 + hLen, maskedDB, 0, dbLen);

        // 计算seedmask
        byte[] seedmask = MGF(maskedDB, hLen);
        // seed = maskedSeed ^ seedmask
        byte[] seed = new byte[hLen];
        for (int i = 0; i < hLen; i++) {
            assert seedmask != null;
            seed[i] = (byte) (maskedSeed[i] ^ seedmask[i]);
        }

        // 计算 DBmask
        byte[] DBmask = MGF(seed, dbLen);
        byte[] DB = new byte[dbLen];
        for (int i = 0; i < dbLen; i++) {
            assert DBmask != null;
            DB[i] = (byte) (maskedDB[i] ^ DBmask[i]);
        }

        int index = -1;
        for (int i = hLen; i < dbLen; i++) {
            if (DB[i] == 0x01) {
                index = i + 1;
                break;
            }
        }
        byte[] plainOAEP = new byte[dbLen - index];
        System.arraycopy(DB, index, plainOAEP, 0, plainOAEP.length);

        return plainOAEP;
    }

    /**
     * 掩码生成函数
     *
     * @param X       消息
     * @param maskLen 掩码长度
     * @return 输出为产生的掩码
     */
    private byte[] MGF(byte[] X, int maskLen) {
        byte[] HM = SHA1.digest(X);
        if (maskLen == HM.length) {
            return HM;
        } else if (maskLen > HM.length) {
            byte[] res = new byte[maskLen];
            System.arraycopy(HM, 0, res, 0, HM.length);
            for (int i = 0; i < maskLen - HM.length; i++) {
                res[HM.length + i] = 0x00;
            }
            return res;
        }
        return null;
    }
}
